# database/connection.py
import mysql.connector
from mysql.connector import Error
from config.settings import settings
from contextlib import contextmanager
from contextlib import contextmanager
from database.models import PushTask
# 创建数据库连接
def create_connection():
    try:
        connection = mysql.connector.connect(
            host=settings.DB_HOST,
            user=settings.DB_USER,
            password=settings.DB_PASSWORD,
            database=settings.DB_NAME,
            port=settings.DB_PORT,
            charset="utf8mb4",
            autocommit=True  # 根据业务需求调整
        )
        return connection
    except Error as e:
        print(f"数据库连接错误: {e}")
        raise

# 连接池实现（可选，用于提高性能）
class DatabasePool:
    _instance = None
    
    def __new__(cls):
        if not cls._instance:
            cls._instance = super().__new__(cls)
            cls._instance.pool = create_connection()
        return cls._instance
    
    def get_connection(self):
        return self.pool

# 连接上下文管理器
@contextmanager
def get_db():
    db = create_connection()
    try:
        yield db
    finally:
        db.close()

def get_pending_tasks(limit: int = 10) -> list[PushTask]:
    # """从队列获取待推送任务"""
    with get_db() as db:
        cursor = db.cursor(dictionary=True)
        query = """
            SELECT a.id, a.tempmsg, b.app_open_id 
            FROM ci_messages a 
            INNER JOIN t_fullmarketing_user_app b 
            ON a.userid = b.uc_uid 
            WHERE a.app_flag = 0 
            LIMIT %s
        """
        cursor.execute(query, (limit,))
        rows = cursor.fetchall()
        task_ids = [row['id'] for row in rows]
        print(f"取走的任务 ID: {task_ids}")
        return [PushTask(**row) for row in rows]

def mark_task_status(task_id: int, status: int) -> bool:
    """更新任务状态
    1: 取走
    2: 执行成功
    3: 执行失败
    """
    with get_db() as db:
        cursor = db.cursor()
        query = "UPDATE ci_messages SET app_flag = %s WHERE id = %s"
        cursor.execute(query, (status, task_id))
        return cursor.rowcount > 0